package token

import (
	"context"
	"crypto/rsa"
	"github.com/dgrijalva/jwt-go"
	"google.golang.org/grpc/metadata"
	"pay/pkg/ecode"
	"strings"
)

type ContextKey struct{}

// Token
type Token struct {
	Type   string   // token类型
	Raw    string   // token本体
	Digest string   // token解密后识别关键字
	Body   MyClaims // token对象
}

type MyClaims interface {
	jwt.Claims
	String() string
}

func GetTokenFromContext(ctx context.Context) *Token {
	token, ok := ctx.Value(ContextKey{}).(*Token)
	if ok {
		return token
	} else {
		return &Token{}
	}
}

// InitRPCToken
func InitRPCToken(ctx context.Context, tokenType, metaKey string, jwtPublicKey *rsa.PublicKey, claims MyClaims) (
	context.Context, error) {
	token := &Token{Type: tokenType}
	if md, ok := metadata.FromIncomingContext(ctx); ok {
		tokens := md.Get(metaKey)
		if len(tokens) == 1 {
			token.Raw = tokens[0]
			if err := extractJwtToken(token.Raw, jwtPublicKey, claims); err != nil {
				return ctx, err
			} else {
				token.Body = claims
				token.Digest = claims.String()
				return context.WithValue(ctx, ContextKey{}, token), nil
			}

		}
	}
	return ctx, ecode.Unauthenticated("登录已失效，请重新登录")
}

func extractJwtToken(rawToken string, jwtPublicKey *rsa.PublicKey, claims MyClaims) error {
	strArr := strings.Split(rawToken, " ")
	var tokenString string
	if len(strArr) == 1 {
		tokenString = rawToken
	} else if len(strArr) == 2 {
		tokenString = strArr[1]
	} else {
		return ecode.Unauthenticated("登录已失效，请重新登录")
	}

	var (
		jwtToken *jwt.Token
		err      error
	)
	if jwtToken, err = new(jwt.Parser).ParseWithClaims(tokenString, claims, func(token *jwt.Token) (interface{}, error) {
		//Make sure that the token method conform to "SigningMethodHMAC"
		if _, ok := token.Method.(*jwt.SigningMethodRSA); !ok {
			return nil, ecode.Unauthenticated("登录已失效，请重新登录")
		}
		return jwtPublicKey, nil
	}); err != nil {
		return ecode.Unauthenticated("登录已失效，请重新登录")
	}
	if !jwtToken.Valid {
		return ecode.Unauthenticated("登录已失效，请重新登录")
	} else {
		return nil
	}
}
